#define _USE_MATH_DEFINES
#include <cmath>
#include "AlgoPRM.h"
#include "AlgoPRMPrivate.h"
#include "DebugConfig.h"
#include <iostream>
#include <fstream>
#include <sstream>
#include <queue>
#include <algorithm>
#include <limits>
#include <set>

namespace algo
{

// AlgoPRM public methods (PIMPL pattern)
AlgoPRM::AlgoPRM() : d(new AlgoPRMPrivate()) {}
AlgoPRM::~AlgoPRM() { delete d; }
int AlgoPRM::buildRoadmap(const algobase::TaskConfig& taskConfig, const std::vector<std::vector<double>>& terrain) {
    return d->buildRoadmap(taskConfig, terrain);
}
int AlgoPRM::findPaths(const algobase::TaskConfig& taskConfig, algobase::Paths& paths) {
    return d->findPaths(taskConfig, paths);
}
int AlgoPRM::saveRoadmap(const std::string& filepath) const {
    return d->saveRoadmap(filepath);
}
int AlgoPRM::loadRoadmap(const std::string& filepath) {
    return d->loadRoadmap(filepath);
}


// AlgoPRMPrivate implementation
AlgoPRMPrivate::AlgoPRMPrivate() {
    random_generator_ = std::make_unique<std::mt19937>(std::random_device{}());
}
AlgoPRMPrivate::~AlgoPRMPrivate() {}

// --- Offline Phase ---
int AlgoPRMPrivate::buildRoadmap(const algobase::TaskConfig& taskConfig, const std::vector<std::vector<double>>& terrain) {
    std::cout << "Building roadmap..." << std::endl;
    current_task_config_ = &taskConfig;
    terrain_data_ = terrain;
    graph_.clear();
    next_node_id_ = 0;

    if (terrain_data_.empty()) return -1;

    std::vector<algobase::CoordZ> sampled_points = sampleNodes();
    if (sampled_points.empty()) return -2;

    for (const auto& pos : sampled_points) {
        int new_id = next_node_id_++;
        graph_.emplace(new_id, RoadmapNode(new_id, pos));
    }

    connectNodes();
    std::cout << "Roadmap built with " << graph_.size() << " nodes." << std::endl;
    checkConnectivity();
    return 0;
}

std::vector<algobase::CoordZ> AlgoPRMPrivate::sampleNodes() {
    std::vector<algobase::CoordZ> points;
    const auto& taskConfig = *current_task_config_;
    int total_samples = 1000; // Configurable

    while(points.size() < total_samples) {
        std::uniform_real_distribution<> dis_lon(taskConfig.mMinLongitude, taskConfig.mMaxLongitude);
        std::uniform_real_distribution<> dis_lat(taskConfig.mMinLatitude, taskConfig.mMaxLatitude);
        algobase::Coord random_coord = {dis_lon(*random_generator_), dis_lat(*random_generator_)};
        
        bool is_valid = isPointValid(random_coord);
        DEBUG_LOG("AlgoPRM.log_sampling", "Sampled point (" << random_coord[0] << ", " << random_coord[1] << "). Valid: " << (is_valid ? "Yes" : "No"));
        if (is_valid) {
            points.push_back({random_coord[0], random_coord[1], getDepth(random_coord)});
        }
    }
    // TODO: Add enhanced sampling for mandatory zones
    return points;
}

void AlgoPRMPrivate::connectNodes() {
    int K = 10; // Connect to K nearest neighbors
    for (auto const& [node_id, node] : graph_) {
        auto neighbors = findKNearestNeighbors(node.position, K);
        for (int neighbor_id : neighbors) {
            if (node_id >= neighbor_id) continue;
            const auto& neighbor_node = graph_.at(neighbor_id);
            bool is_valid = isEdgeValid(node.position, neighbor_node.position);
            DEBUG_LOG("AlgoPRM.log_connectivity", "Checking edge between node " << node_id << " and " << neighbor_id << ". Valid: " << (is_valid ? "Yes" : "No"));
            if (is_valid) {
                double cost = distance(node.position, neighbor_node.position);
                graph_[node_id].neighbors.push_back({neighbor_id, cost});
                graph_[neighbor_id].neighbors.push_back({node_id, cost});
            }
        }
    }
}

void AlgoPRMPrivate::checkConnectivity() const {
    if (graph_.empty()) {
        std::cout << "Connectivity Check: Graph is empty." << std::endl;
        return;
    }

    std::set<int> visited;
    std::vector<int> stack;
    int components = 0;

    for (const auto& pair : graph_) {
        int start_node_id = pair.first;
        if (visited.find(start_node_id) == visited.end()) {
            components++;
            stack.push_back(start_node_id);
            visited.insert(start_node_id);

            while (!stack.empty()) {
                int current_id = stack.back();
                stack.pop_back();

                if (graph_.count(current_id)) {
                    for (const auto& neighbor_pair : graph_.at(current_id).neighbors) {
                        int neighbor_id = neighbor_pair.first;
                        if (visited.find(neighbor_id) == visited.end()) {
                            visited.insert(neighbor_id);
                            stack.push_back(neighbor_id);
                        }
                    }
                }
            }
        }
    }

    std::cout << "Connectivity Check: Found " << components << " connected component(s) in the graph." << std::endl;
    if (components > 1) {
        std::cout << "Warning: Roadmap graph is not fully connected. Pathfinding may fail between different components." << std::endl;
    } else {
        std::cout << "Info: Roadmap graph is fully connected." << std::endl;
    }
}

// --- Online Phase ---
int AlgoPRMPrivate::findPaths(const algobase::TaskConfig& taskConfig, algobase::Paths& paths) {
    current_task_config_ = &taskConfig; // Update task config for this query
    paths.clear();

    int start_node_id = connectPointToGraph(taskConfig.mStartPos, "start");
    int end_node_id = connectPointToGraph(taskConfig.mEndPos, "end");
    if (start_node_id == -1 || end_node_id == -1) return -1;

    std::vector<std::vector<int>> waypoint_groups;
    const auto mandatory_polygons_it = taskConfig.mTypePolygons.find(algobase::AT_Mandatory);
    if (mandatory_polygons_it != taskConfig.mTypePolygons.end()) {
        for (const auto& polygon : mandatory_polygons_it->second) {
            waypoint_groups.push_back(findRepresentativeNodeIdsForZone(polygon));
        }
    }

    std::vector<int> from_nodes = {start_node_id};
    std::vector<std::vector<algobase::Path>> all_segment_paths;

    for(const auto& group : waypoint_groups) {
        // For simplicity, we find paths from all previous endpoints to the first node of the current mandatory zone
        // A more complex implementation would find the optimal connection.
        int to_node = group.empty() ? -1 : group[0];
        if (to_node == -1) continue;

        std::vector<algobase::Path> segment_paths;
        for(int from_node : from_nodes) {
            auto found_paths = findKShortestPaths(from_node, to_node, 2); // Find 2 paths for diversity
            segment_paths.insert(segment_paths.end(), found_paths.begin(), found_paths.end());
        }
        all_segment_paths.push_back(segment_paths);
        from_nodes = {to_node};
    }
    
    std::vector<algobase::Path> final_segment_paths;
    for(int from_node : from_nodes) {
        auto found_paths = findKShortestPaths(from_node, end_node_id, 5);
        final_segment_paths.insert(final_segment_paths.end(), found_paths.begin(), found_paths.end());
    }
    all_segment_paths.push_back(final_segment_paths);

    paths = combinePaths(all_segment_paths);

    disconnectPointFromGraph(start_node_id);
    disconnectPointFromGraph(end_node_id);
    return 0;
}

int AlgoPRMPrivate::connectPointToGraph(const algobase::Coord& point, const std::string& type) {
    if (!isPointValid(point)) return -1;
    
    int new_id = (type == "start") ? -1 : -2; // Use negative IDs for temporary nodes
    algobase::CoordZ pos = {point[0], point[1], getDepth(point)};
    graph_.emplace(new_id, RoadmapNode(new_id, pos));

    int K = 10;
    auto neighbors = findKNearestNeighbors(pos, K);
    for (int neighbor_id : neighbors) {
        if (isEdgeValid(pos, graph_.at(neighbor_id).position)) {
            double cost = distance(pos, graph_.at(neighbor_id).position);
            graph_[new_id].neighbors.push_back({neighbor_id, cost});
            graph_[neighbor_id].neighbors.push_back({new_id, cost});
        }
    }
    return new_id;
}

void AlgoPRMPrivate::disconnectPointFromGraph(int node_id) {
    if (graph_.count(node_id)) {
        for (const auto& neighbor_pair : graph_.at(node_id).neighbors) {
            auto& neighbor_node = graph_.at(neighbor_pair.first);
            neighbor_node.neighbors.erase(
                std::remove_if(neighbor_node.neighbors.begin(), neighbor_node.neighbors.end(),
                               [node_id](const auto& p){ return p.first == node_id; }),
                neighbor_node.neighbors.end()
            );
        }
        graph_.erase(node_id);
    }
}

std::vector<int> AlgoPRMPrivate::findRepresentativeNodeIdsForZone(const algobase::Polygon& polygon) {
    std::vector<int> nodes_in_zone;
    for(const auto& pair : graph_) {
        if(pair.first < 0) continue; // Skip temp nodes
        algobase::Coord p2d = {pair.second.position[0], pair.second.position[1]};
        if(isInside(p2d, polygon)) {
            nodes_in_zone.push_back(pair.first);
        }
    }
    return nodes_in_zone;
}

// Basic Dijkstra for finding the single shortest path (used by Yen's)
std::vector<int> dijkstra(int start_id, int end_id, const std::map<int, RoadmapNode>& graph) {
    std::priority_queue<std::pair<double, int>, std::vector<std::pair<double, int>>, std::greater<std::pair<double, int>>> pq;
    std::map<int, double> dist;
    std::map<int, int> prev;

    for (const auto& pair : graph) {
        dist[pair.first] = std::numeric_limits<double>::infinity();
    }

    dist[start_id] = 0;
    pq.push({0, start_id});

    while (!pq.empty()) {
        double d = pq.top().first;
        int u = pq.top().second;
        pq.pop();

        if (u == end_id) break;
        if (d > dist[u]) continue;

        for (const auto& neighbor_pair : graph.at(u).neighbors) {
            int v = neighbor_pair.first;
            double weight = neighbor_pair.second;
            if (dist[u] + weight < dist[v]) {
                dist[v] = dist[u] + weight;
                prev[v] = u;
                pq.push({dist[v], v});
            }
        }
    }

    std::vector<int> path;
    if (prev.count(end_id) || start_id == end_id) {
        int curr = end_id;
        while (curr != start_id) {
            path.push_back(curr);
            if(prev.find(curr) == prev.end()) return {}; // Path broken
            curr = prev[curr];
        }
        path.push_back(start_id);
        std::reverse(path.begin(), path.end());
    }
    return path;
}


std::vector<algobase::Path> AlgoPRMPrivate::findKShortestPaths(int start_id, int end_id, int K) {
    std::vector<std::vector<int>> A; // Stores the K shortest paths found so far
    std::priority_queue<std::pair<double, std::vector<int>>, std::vector<std::pair<double, std::vector<int>>>, std::greater<std::pair<double, std::vector<int>>>> B; // Candidate paths

    // Find the first shortest path
    auto first_path = dijkstra(start_id, end_id, graph_);
    if (first_path.empty()) {
        DEBUG_LOG("AlgoPRM.log_path_search", "Dijkstra failed to find any path between " << start_id << " and " << end_id);
        return {};
    }
    A.push_back(first_path);

    for (int k = 1; k < K; ++k) {
        const auto& previous_path = A[k - 1];

        for (size_t i = 0; i < previous_path.size() - 1; ++i) {
            int spur_node_id = previous_path[i];
            std::vector<int> root_path(previous_path.begin(), previous_path.begin() + i + 1);

            std::map<int, RoadmapNode> temp_graph = graph_;
            
            // Remove edges that are part of the previous shortest paths
            for (const auto& path_in_A : A) {
                if (path_in_A.size() > i) {
                    bool same_root = true;
                    for(size_t j = 0; j <= i; ++j) {
                        if(path_in_A[j] != root_path[j]) {
                            same_root = false;
                            break;
                        }
                    }
                    if (same_root) {
                        int u = path_in_A[i];
                        int v = path_in_A[i+1];
                        // Remove edge u-v from temp_graph
                        temp_graph[u].neighbors.erase(
                            std::remove_if(temp_graph[u].neighbors.begin(), temp_graph[u].neighbors.end(),
                                           [v](const auto& p){ return p.first == v; }),
                            temp_graph[u].neighbors.end()
                        );
                    }
                }
            }

            // Remove nodes on the root path (except the spur node) from the temp graph
            for (size_t j = 0; j < i; ++j) {
                temp_graph.erase(root_path[j]);
            }

            auto spur_path = dijkstra(spur_node_id, end_id, temp_graph);

            if (!spur_path.empty()) {
                std::vector<int> total_path = root_path;
                total_path.insert(total_path.end(), spur_path.begin() + 1, spur_path.end());
                
                double total_cost = 0;
                for(size_t j=0; j < total_path.size() - 1; ++j) {
                    total_cost += distance(graph_.at(total_path[j]).position, graph_.at(total_path[j+1]).position);
                }
                B.push({total_cost, total_path});
            }
        }

        if (B.empty()) {
            break; // No more paths
        }

        // Add the best candidate path to A
        bool new_path_found = false;
        while(!B.empty()) {
            auto candidate = B.top();
            B.pop();
            bool is_duplicate = false;
            for(const auto& path_in_A : A) {
                if(path_in_A == candidate.second) {
                    is_duplicate = true;
                    break;
                }
            }
            if(!is_duplicate) {
                A.push_back(candidate.second);
                new_path_found = true;
                break;
            }
        }
        if(!new_path_found) break; // No more unique paths
    }

    std::vector<algobase::Path> result_paths;
    for (const auto& path_ids : A) {
        result_paths.push_back(reconstructGraphPath(path_ids));
    }
    return result_paths;
}

void combineRecursive(
    size_t segment_index,
    algobase::Path current_path,
    const std::vector<std::vector<algobase::Path>>& all_segment_paths,
    algobase::Paths& final_paths,
    size_t& path_count_limit
) {
    if (final_paths.size() >= path_count_limit) {
        return;
    }
    if (segment_index == all_segment_paths.size()) {
        final_paths.push_back(current_path);
        return;
    }

    const auto& next_segment_options = all_segment_paths[segment_index];
    for (const auto& segment : next_segment_options) {
        if (segment.empty()) continue;
        algobase::Path next_path = current_path;
        // Append the next segment, skipping the first point to avoid duplication
        next_path.insert(next_path.end(), segment.begin() + 1, segment.end());
        combineRecursive(segment_index + 1, next_path, all_segment_paths, final_paths, path_count_limit);
        if (final_paths.size() >= path_count_limit) {
            return;
        }
    }
}

algobase::Paths AlgoPRMPrivate::combinePaths(const std::vector<std::vector<algobase::Path>>& all_segment_paths) {
    algobase::Paths final_paths;
    if (all_segment_paths.empty()) return {};

    size_t path_limit = 5000; // Set a limit to avoid generating millions of paths

    // Start the recursion for each path in the first segment
    for (const auto& first_segment : all_segment_paths[0]) {
        combineRecursive(1, first_segment, all_segment_paths, final_paths, path_limit);
        if (final_paths.size() >= path_limit) {
            break;
        }
    }
    
    return final_paths;
}

algobase::Path AlgoPRMPrivate::reconstructGraphPath(const std::vector<int>& node_ids) {
    algobase::Path path;
    for(int id : node_ids) {
        algobase::PathCoord pc;
        pc.mCoordZ = graph_.at(id).position;
        path.push_back(pc);
    }
    return path;
}


// --- Common Helpers ---
// (Implementations for isPointValid, isEdgeValid, getDepth, isInside, distance, findKNearestNeighbors)
// These would be similar to the implementations in AlgoRRTStar.cpp
bool AlgoPRMPrivate::isPointValid(const algobase::Coord& p) const {
    if (p[0] < current_task_config_->mMinLongitude || p[0] > current_task_config_->mMaxLongitude ||
        p[1] < current_task_config_->mMinLatitude  || p[1] > current_task_config_->mMaxLatitude) {
        return false;
    }
    if (getDepth(p) > 0) return false; // On land
    for (const auto& pair : current_task_config_->mTypePolygons) {
        if (pair.first == algobase::AT_Obstacle || pair.first == algobase::AT_Restricted) {
            for (const auto& polygon : pair.second) {
                if (isInside(p, polygon)) return false;
            }
        }
    }
    return true;
}

bool AlgoPRMPrivate::isEdgeValid(const algobase::CoordZ& p1, const algobase::CoordZ& p2) const {
    int steps = 10;
    for (int i = 0; i <= steps; ++i) {
        double t = static_cast<double>(i) / steps;
        algobase::Coord p_2d = { p1[0] + t * (p2[0] - p1[0]), p1[1] + t * (p2[1] - p1[1]) };
        if (!isPointValid(p_2d)) return false;
    }
    return true;
}

double AlgoPRMPrivate::getDepth(const algobase::Coord& coord) const {
    if (terrain_data_.empty() || terrain_data_[0].empty()) return 0.0;
    double lon_frac = (coord[0] - current_task_config_->mMinLongitude) / (current_task_config_->mMaxLongitude - current_task_config_->mMinLongitude);
    double lat_frac = (coord[1] - current_task_config_->mMinLatitude) / (current_task_config_->mMaxLatitude - current_task_config_->mMinLatitude);
    int x = static_cast<int>(lon_frac * (terrain_data_[0].size() - 1));
    int y = static_cast<int>(lat_frac * (terrain_data_.size() - 1));
    x = std::max(0, std::min((int)terrain_data_[0].size() - 1, x));
    y = std::max(0, std::min((int)terrain_data_.size() - 1, y));
    return terrain_data_[y][x];
}

bool AlgoPRMPrivate::isInside(const algobase::Coord& point, const algobase::Polygon& polygon) const {
    if (polygon.size() < 3) return false;
    int crossings = 0;
    for (size_t i = 0; i < polygon.size(); ++i) {
        const auto& p1 = polygon[i];
        const auto& p2 = polygon[(i + 1) % polygon.size()];
        if (((p1[1] <= point[1] && point[1] < p2[1]) || (p2[1] <= point[1] && point[1] < p1[1])) &&
            (point[0] < (p2[0] - p1[0]) * (point[1] - p1[1]) / (p2[1] - p1[1]) + p1[0])) {
            crossings++;
        }
    }
    return (crossings % 2) != 0;
}

double AlgoPRMPrivate::distance(const algobase::CoordZ& p1, const algobase::CoordZ& p2) const {
    double lon1 = p1[0], lat1 = p1[1], depth1 = p1[2];
    double lon2 = p2[0], lat2 = p2[1], depth2 = p2[2];
    double R = 6371e3;
    double x = (lon2 - lon1) * M_PI/180.0 * cos((lat1 + lat2)/2.0 * M_PI/180.0);
    double y = (lat2 - lat1) * M_PI/180.0;
    double horiz_dist = sqrt(x*x + y*y) * R;
    double vert_dist = depth2 - depth1;
    return sqrt(horiz_dist*horiz_dist + vert_dist*vert_dist);
}

std::vector<int> AlgoPRMPrivate::findKNearestNeighbors(const algobase::CoordZ& pos, int K) {
    std::vector<std::pair<double, int>> dists;
    for (const auto& pair : graph_) {
        if (pair.first < 0) continue; // Skip temp nodes
        dists.push_back({distance(pos, pair.second.position), pair.first});
    }
    std::sort(dists.begin(), dists.end());
    std::vector<int> result;
    for (int i = 0; i < std::min((int)dists.size(), K); ++i) {
        result.push_back(dists[i].second);
    }
    return result;
}

// --- Serialization ---
int AlgoPRMPrivate::saveRoadmap(const std::string& filepath) const {
    std::ofstream ofs(filepath, std::ios::binary);
    if (!ofs) return -1;
    size_t graph_size = graph_.size();
    ofs.write(reinterpret_cast<const char*>(&graph_size), sizeof(graph_size));
    for (const auto& pair : graph_) {
        ofs.write(reinterpret_cast<const char*>(&pair.second.id), sizeof(int));
        ofs.write(reinterpret_cast<const char*>(&pair.second.position), sizeof(algobase::CoordZ));
        size_t neighbors_size = pair.second.neighbors.size();
        ofs.write(reinterpret_cast<const char*>(&neighbors_size), sizeof(neighbors_size));
        ofs.write(reinterpret_cast<const char*>(pair.second.neighbors.data()), neighbors_size * sizeof(std::pair<int, double>));
    }
    return 0;
}

int AlgoPRMPrivate::loadRoadmap(const std::string& filepath) {
    std::ifstream ifs(filepath, std::ios::binary);
    if (!ifs) return -1;
    graph_.clear();
    size_t graph_size;
    ifs.read(reinterpret_cast<char*>(&graph_size), sizeof(graph_size));
    for (size_t i = 0; i < graph_size; ++i) {
        int id;
        algobase::CoordZ pos;
        ifs.read(reinterpret_cast<char*>(&id), sizeof(int));
        ifs.read(reinterpret_cast<char*>(&pos), sizeof(algobase::CoordZ));
        RoadmapNode node(id, pos);
        size_t neighbors_size;
        ifs.read(reinterpret_cast<char*>(&neighbors_size), sizeof(neighbors_size));
        node.neighbors.resize(neighbors_size);
        ifs.read(reinterpret_cast<char*>(node.neighbors.data()), neighbors_size * sizeof(std::pair<int, double>));
        graph_[id] = node;
    }
    next_node_id_ = graph_.size();
    return 0;
}

} // namespace algo